# coding=utf-8
from flask import Flask
from flask import request, jsonify, \
    send_from_directory
import string, json, base64
import tensorflow as tf
import vis_lstm_model
import data_loader
import numpy as np
import utils
import re, os


#################################################################
########################  Parameters  ###########################

model_path = './Data/model2.ckpt'
num_lstm_layers = 2
fc7_feature_length = 4096
rnn_size = 512
embedding_size = 512
word_emb_dropout = 0.5
image_dropout = 0.5
data_dir = 'Data'

vocab_data = data_loader.get_question_answer_vocab(data_dir)
question_vocab = vocab_data['question_vocab']
question_ids = np.zeros((1, vocab_data['max_question_length']), dtype = 'int32')
vgg_graph, fc7_tensor, images = utils.get_vgg_graph(os.path.join(data_dir, 'vgg16.tfmodel'))

model_options = {
	'num_lstm_layers' : num_lstm_layers,
	'rnn_size' : rnn_size,
	'embedding_size' : embedding_size,
	'word_emb_dropout' : word_emb_dropout,
	'image_dropout' : image_dropout,
	'fc7_feature_length' : fc7_feature_length,
	'lstm_steps' : vocab_data['max_question_length'] + 1,
	'q_vocab_size' : len(vocab_data['question_vocab']),
	'ans_vocab_size' : len(vocab_data['answer_vocab'])
}
model = vis_lstm_model.Vis_lstm_model(model_options)
with model.graph.as_default():
	input_tensors, t_prediction, t_ans_probab = model.build_generator()
tf.reset_default_graph()



#################################################################
########################  Prediction  ###########################

def predict(img, question='Which animal is this?'):
	print "\n****** Prediction ******\n"

	fc7_features = utils.get_fc7_feas(img, vgg_graph, fc7_tensor, images)

	word_regex = re.compile(r'\w+')
	question_words = re.findall(word_regex, question)
	base = vocab_data['max_question_length'] - len(question_words)
	for i in range(0, len(question_words)):
		if question_words[i] in question_vocab:
			question_ids[0][base + i] = question_vocab[ question_words[i] ]
		else:
			question_ids[0][base + i] = question_vocab['UNK']

	with tf.Session(graph = model.graph) as sess:
		saver = tf.train.Saver()
		saver.restore(sess, model_path)
		# do predition
		pred, answer_probab = sess.run([t_prediction, t_ans_probab], feed_dict={
			input_tensors['fc7']:fc7_features,
			input_tensors['sentence']:question_ids,
		})

	ans_map = { vocab_data['answer_vocab'][ans] : ans for ans in vocab_data['answer_vocab']}
	print "Ans:", ans_map[pred[0]]

	answer_probab_tuples = [(-answer_probab[0][idx], idx) for idx in range(len(answer_probab[0]))]
	answer_probab_tuples.sort()
	print "Top Answers:"
	ret = []
	for i in range(5):
		print i, ":", ans_map[ answer_probab_tuples[i][1] ]
		ret.append(ans_map[ answer_probab_tuples[i][1] ])

	return ret



#################################################################
##########################  Flask  ##############################

app = Flask(__name__, static_url_path='')

class InvalidUsage(Exception):
	status_code = 400

	def __init__(self, message, status_code=None, payload=None):
		Exception.__init__(self)
		self.message = message
		if status_code is not None:
			self.status_code = status_code
		self.payload = payload

	def to_dict(self):
		rv = dict(self.payload or ())
		rv['message'] = self.message
		return rv

@app.errorhandler(InvalidUsage)
def handle_invalid_usage(error):
	response = jsonify(error.to_dict())
	response.status_code = error.status_code
	return response


# SEARCH
@app.route("/get_tags", methods=['GET', 'POST'])
def get_tags():
	"""get tags corresponding to a image"""
	if not 'img' in request.files:
		raise InvalidUsage('parameter "img" is missing', status_code=410)
	if not 'question' in request.form:
		raise InvalidUsage('parameter "question" is missing', status_code=410)
	try:
		f = request.files.get('img')
		img_str = f.read()
		img = utils.read_img_blob(img_str)
		que_str = request.form.get("question")
		print "question:",que_str
	except:
		raise InvalidUsage('Invalid "img" param, must be a blob string',
							status_code=410)
	out = {}
	out['tags'] = predict(img, question=que_str)

	return jsonify(out)


@app.route('/static/<path:path>')
def send_static_files(path):
	"static files"
	return send_from_directory('static_data', path)

@app.route('/img/<path:path>')
def send_image(path):
	"static files"
	return send_from_directory(IMGS_PATH, path)


if __name__ == '__main__':
	app.run(host='0.0.0.0', port=5000, debug=True)
